import sys

sys.path.append(".")

import argparse
from pathlib import Path

import onnxruntime as ort

from app.process import DataSource, run_process_image, run_process_video
from src.models.decoder import Decoder


def main(args):
    inputs = args.inputs
    weight = args.weights
    output_dir = args.outputs

    # 初始化数据源
    data_source = DataSource(inputs)
    # 初始化模型和解码器
    model = ort.InferenceSession(weight, providers=["CUDAExecutionProvider"])
    decoder = Decoder(0.25, scale=64)

    # output path
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # 控制参数
    cfg = {
        "output": output_path,
        "save_mode": args.save_mode,
        "scale_g": args.scale,
        "interval": args.interval,
        "show": args.show,
        "sizes": input_size,
        "conf_thres": args.conf_thres,
        "max_age": args.max_age,
        "min_hit": args.min_hit,
        "iou_thres": args.iou_thres,
        "min_pixel": args.min_pixel,
        "period": args.period,
    }

    if data_source.mode == "video":
        run_process_video(data_source, model, decoder, cfg)
    elif data_source.mode == "image":
        run_process_image(data_source, model, decoder, cfg)
    return


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--inputs", type=str)
    parser.add_argument("--weights", type=str)
    parser.add_argument("--outputs", default="resources\\demo", type=str)
    parser.add_argument("--scale", default=0.5, type=float)
    parser.add_argument("--show", action="store_true")
    parser.add_argument("--save_mode", default="image", type=str)
    parser.add_argument("--interval", default=4, type=int)
    parser.add_argument("--conf_thres", default=0.25, type=float)
    parser.add_argument("--max_age", default=1, type=int)
    parser.add_argument("--min_hit", default=3, type=int)
    parser.add_argument("--iou_thres", default=0.3, type=float)
    parser.add_argument("--min_pixel", default=10, type=int)
    parser.add_argument("--period", default=None, type=str, help="Example: 0:15-3:20")

    args = parser.parse_args()

    input_size = [(512, 288), (1024, 576)]

    main(args)
